import re, time, traceback, zipfile
from pathlib import Path
from modules.common.MTKLog import *
from modules.common.Logger import Logger
from modules.common.Device import get_device_info
from modules.common.GlobalAttrs import *
from openpyxl import Workbook
from openpyxl.utils import get_column_letter
import os

class AppStartFactory():
    start_opt_dict = {0:" -W ", 
     1:" -n ",  2:" "}

    def __init__(self, sn, executeView):
        self.sn = sn
        self.atrace = executeView.atrace
        self.async_dump = executeView.async_dump
        self.atrace_duration = executeView.atrace_duration
        self.atrace_options = executeView.atrace_options
        self.atrace_options_async = executeView.atrace_options_async
        self.package_info_list = executeView.package_info_list
        self.total_round = executeView.total_round
        self.start_opt = executeView.start_opt
        self.clean_before = executeView.clean_before
        self.clean = executeView.clean
        self.app_start_remain = executeView.app_start_remain
        self.surfaceFlinger = executeView.surfaceFlinger
        self.log_cycle_type = executeView.log_cycle_type
        self.memdump_type = executeView.memdump_type
        self.dumpsys_type = executeView.dumpsys_type
        self.cmd_list_before_apk = executeView.cmd_list_before_apk
        self.cmd_list_after_apk = executeView.cmd_list_after_apk
        self.cmd_list_after_round = executeView.cmd_list_after_round
        self.screencap = executeView.screencap
        self.wait_before_screencap = executeView.wait_before_screencap
        self.log_path = os.path.join(executeView.result_timestamp_dir, f"AppStart_{sn}.log")
        self.result_dir = os.path.join(executeView.result_timestamp_dir, sn)
        self.root_dir = executeView.result_root
        self.startTime = int(time.time())
        self.endTime = int(time.time())
        self.build = None
        self.project = None
        self.screencap_dir = None
        self.atrace_dir = os.path.join(self.result_dir, "atrace")
        if not os.path.isdir(self.atrace_dir):
            os.makedirs(self.atrace_dir)
        if self.screencap:
            self.screencap_dir = os.path.join(self.result_dir, "screencap")
            if not os.path.isdir(self.screencap_dir):
                os.makedirs(self.screencap_dir)
        self.segcreate = None
        self.end_test = False
        self.segCreateTest = executeView.segCreateTest
        self.round_times_test = executeView.round_times_test
        self.old_test_months = executeView.old_test_months
        self.fragment_percent_test = executeView.fragment_percent_test
        self.size_copy_test = executeView.size_copy_test
        self.time_to_delay = executeView.time_to_delay
        self.clear_for_data = executeView.clear_for_data
    def save_to_excel(self, result_data, file_path):
        wb = Workbook()
        
        for sn, apk_data in result_data.items():
            ws = wb.create_sheet(title=sn)

            # 写入表头
            headers = ["测试 APK 名称"] + [f"第 {i} 轮启动时间 (ms)" for i in range(1, 9)]
            ws.append(headers)

            # 写入每个 APK 的启动时间
            for package_name, times in apk_data.items():
                # 确保每个应用的启动时间列表长度为 8，如果不足则填充为 0
                while len(times) < 8:
                    times.append(0)
                row_data = [package_name] + [int(time) for time in times]
                ws.append(row_data)

            # 调整列宽
            for col_idx in range(1, len(headers) + 1):
                col_letter = get_column_letter(col_idx)
                ws.column_dimensions[col_letter].width = 20

        # 删除默认的Sheet（openpyxl 默认会生成一个名为“Sheet”的 Sheet）
        if "Sheet" in wb.sheetnames:
            del wb["Sheet"]

        # 保存 Excel 文件
        wb.save(file_path)
        print(f"测试结果已保存到 {file_path}")

    def run(self):
        logger = Logger(self.log_path)
        logger.info(f"SerialNumber：{self.sn} --- APP启动测试开始\n")
        logger.info(f"SerialNumber：{self.sn} --- 获取手机信息\n")
        # 仅为保持与原实现一致：调用获取信息并打印完成日志
        _ = get_device_info(self.sn)
        deviceInfoTxt = os.path.join(self.result_dir, "deviceInfo.txt")
        logger.info(f"SerialNumber：{self.sn} --- 读取手机信息完成：{deviceInfoTxt}\n")

        logger.info(f"SerialNumber：{self.sn} --- 开始APP轮转测试，计划测试{self.total_round}轮\n")
        logger.info(f"SerialNumber：{self.sn} --- 启动参数值：{AppStartFactory.start_opt_dict[self.start_opt]}\n")

        # 设备保持常亮
        logger.info(f"SerialNumber：{self.sn} --- 执行命令：adb shell svc power stayon true")
        execute_cmd(f"adb -s {self.sn} shell svc power stayon true")
        
        # 结果与状态容器（仅当前设备，但结构保留多设备扩展性）
        result_data = {self.sn: {}}
        cold_start_times_dict = {self.sn: []}  # 每轮冷启动总时长
        test_status_dict = {self.sn: {"status": "running", "total_time": 0, "average_time": 0}}

        # 轮转测试
        for round_idx in range(1, self.total_round + 1):
            logger.info(f"SerialNumber: {self.sn} --- Round:{round_idx} --- 开始测试\n")
            round_total_ms = 0

            for package_info in self.package_info_list:
                if self.clean_before:
                    logger.info("SerialNumber: " + self.sn + " --- Round:" + str(round_idx) + " --- 停止 " + package_info.package_name)
                    cmdline = "adb -s " + self.sn + " shell am force-stop " + package_info.package_name
                execute_cmd_infinite(cmdline)
                package_name = package_info.package_name
                activity_info_list = package_info.activity_list
                if len(activity_info_list) > 0:
                    wait_time_ms = 0
                    delta_time = 0
                    cmdline = "adb -s " + str(self.sn) + " shell sync"
                    execute_cmd(cmdline)
                # 初始化包结果列表
                if package_name not in result_data[self.sn]:
                    result_data[self.sn][package_name] = []

                # 可选清理
                if self.clean_before:
                    logger.info(f"SerialNumber: {self.sn} --- Round:{round_idx} --- 停止 {package_name}")
                    execute_cmd_infinite(f"adb -s {self.sn} shell am force-stop {package_name}")
                    time.sleep(2)

                logger.info(f"SerialNumber: {self.sn} --- Round:{round_idx} --- 启动 {package_name}")

                activities = getattr(package_info, "activity_list", []) or []
                if not activities:
                    logger.warn(f"SerialNumber: {self.sn} --- Round:{round_idx} --- {package_name} 未配置启动 Activity，跳过")
                    result_data[self.sn][package_name].append(0)
                    continue
                if self.atrace:
                    # 停止上一次的记录（如果有），这会清空缓冲区，是关键一步！
                    logger.info(f"SerialNumber: {self.sn} --- Round:{round_idx} --- Stopping any previous atrace session.")
                    stop_cmd = f"adb -s {self.sn} shell atrace {self.atrace_options_async} --async_stop"
                    execute_cmd(stop_cmd)
                    time.sleep(0.5) # 短暂等待确保完全停止
            
                    # 开始一次全新的异步记录
                    logger.info(f"SerialNumber: {self.sn} --- Round:{round_idx} --- Starting new atrace session for {package_name}.")
                    start_cmd = f"adb -s {self.sn} shell atrace -t {self.atrace_duration} {self.atrace_options_async} --async_start"
                    execute_cmd(start_cmd)
                    
                    # 等待1秒，以满足您“抓取启动之前1s左右”的需求
                    time.sleep(1)
                
                # 发起启动并解析 WaitTime（ms）
                elapsed_s = self._launch_and_measure(logger, package_name, activities[0].activity_name)
                wait_time_ms = elapsed_s.get("wait_time_ms", 0)
                
                if wait_time_ms == 0:
                    logger.error("=============================================================================================")
                    logger.error(f"SerialNumber: {self.sn} --- Round:{round_idx} --- {package_name} wait_time_ms == 0")
                    logger.error(f"SerialNumber: {self.sn} --- Detail start info: {elapsed_s.get('error', 'Unknown error')}")
                    logger.error("=============================================================================================\n")
                
                if self.atrace and wait_time_ms > 0:
                    atrace_file = os.path.join(self.atrace_dir, f"{package_info.package_name}_atrace_{str(round_idx).zfill(4)}_{str(wait_time_ms)}")
                    time.sleep(2)
                    logger.info(f"SerialNumber: {self.sn} --- Round:{round_idx} --- 开始执行 atrace")
                    atrace_cmdline_async = f"adb -s {self.sn} shell atrace -t {self.atrace_duration} {self.atrace_options_async} --async_dump"
                    success, error = cmd_exec_and_save(atrace_cmdline_async, atrace_file, timeout=30)
                    if success:
                        atrace_zip = atrace_file + ".zip"
                        try:
                            with zipfile.ZipFile(atrace_zip, "w", zipfile.ZIP_DEFLATED) as zp:
                                zp.write(atrace_file, "." + os.sep + os.path.basename(atrace_file))
                            self._safe_remove_file(atrace_file)
                        except Exception as e:
                            logger.error(f"创建 atrace zip 文件失败: {e}")
                            self._safe_remove_file(atrace_file)
                    else:
                        logger.error(f"atrace 执行失败: {error}")
                        self._safe_remove_file(atrace_file)
                

                # 汇总单轮
                try:
                    round_total_ms += int(wait_time_ms)
                except (ValueError, TypeError):
                    logger.error(f"wait_time_ms 转换失败: {wait_time_ms}")
                    round_total_ms += 0
                
                result_data[self.sn][package_name].append(wait_time_ms)

                # 保障应用间最小停顿
                self._respect_min_interval(logger, round_idx, package_name, elapsed_s.get("elapsed_s", 0))
                
                # 强制停止应用
                force_stop_cmd = ["adb", "-s", self.sn, "shell", "am", "force-stop", package_info.package_name]
                exec_cmd(force_stop_cmd)

            # 记录该轮汇总
            cold_start_times_dict[self.sn].append(round_total_ms)
            
            # 仅在特定条件下重启ADB服务器，避免不必要的重启
            logger.info(f"SerialNumber: {self.sn} --- Round:{round_idx} --- 执行 adb kill-server")
            execute_cmd("adb kill-server")
            time.sleep(2)  # 等待ADB服务器重启
            
            # 轮次收尾命令
            self._run_post_round_cmds(logger, round_idx)

            logger.info(f"SerialNumber: {self.sn} --- Round:{round_idx} --- 结束测试\n\n")
            if self.end_test:
                logger.warn(f"SerialNumber: {self.sn} --- Round:{round_idx} --- end_test为True，跳出Round循环\n")
                break

        # 结束与汇总
        self.endTime = int(time.time())
        sum_list = cold_start_times_dict[self.sn]
        logger.info(f"SerialNumber：{self.sn} --- 每轮冷启动总时长列表：")
        logger.info(f"SerialNumber：{self.sn} --- {', '.join(map(str, sum_list))}")

        avg_time = 0
        if len(sum_list) > 1:
            # 与原实现一致：跳过首轮，计算“后 7 轮平均值”（实际按剩余轮数求均值）
            avg_time = sum(sum_list[1:]) / len(sum_list[1:])
            logger.info(f"SerialNumber：{self.sn} --- APP启动测试结束，后 7 轮冷启动时长平均值为 {avg_time:.2f}ms")

        test_status_dict[self.sn]["status"] = "completed"
        test_status_dict[self.sn]["total_time"] = int(sum(sum_list))
        test_status_dict[self.sn]["average_time"] = avg_time
        if self.atrace:
            if self.async_dump:
                logger.info("SerialNumber: " + self.sn + " --- async_start 模式，关闭atrace")
                cmdline = "adb -s " + self.sn + " shell atrace -t " + str(self.atrace_duration) + " " + self.atrace_options_async + " --async_stop"
                execute_cmd(cmdline)
        logger.info(f"SerialNumber：{self.sn} --- APP启动测试结束，共测试{self.total_round}轮")
        excel_path = os.path.join(self.root_dir, f"app_start_test_results{self.sn}.xlsx")
        self.save_to_excel(result_data, excel_path)

        # 若未来扩展多设备，这里统一打印最终结果
        if all(s["status"] == "completed" for s in test_status_dict.values()):
            logger.info("所有设备测试已完成，打印最终结果：")
            for _sn, status in test_status_dict.items():
                logger.info(
                    "SerialNumber：{} --- APP启动测试结束，后 7 轮冷启动时长平均值为 {:.2f}ms"
                    .format(_sn, status["average_time"])
                )


    # ===== 新增的私有辅助方法 =====
    def _run_pre_apk_cmds(self, logger, package_name: str, round_idx: int):
        """执行启动前命令并保存输出（与原命名/路径规则保持一致）"""
        for cmd_info in getattr(self, "cmd_list_before_apk", []):
            if "###" not in cmd_info:
                continue
            shell_cmd, rlt_name = cmd_info.split("###", 1)
            cmdline = f"adb -s {self.sn} shell {shell_cmd}"
            rlt_file_path = os.path.join(
                self.result_dir,
                f"{package_name}_{rlt_name}_{str(round_idx).zfill(4)}_0"
            )
            logger.info(f"SerialNumber: {self.sn} --- Round:{round_idx}_0 --- 开始执行: {cmdline}")
            cmd_exec_and_save(cmdline, rlt_file_path, timeout=60)


    def _launch_and_measure(self, logger, package_name: str, activity_name: str):
        """
        启动应用并解析 WaitTime（毫秒）。
        返回: {"wait_time_ms": int, "elapsed_s": int, "error": str}
        """
        # 先清空 SurfaceFlinger 再启动（如有配置）
        if self.surfaceFlinger:
            self.clear_flinger()

        start_time = time.time()
        wait_time_ms = 0
        error_msg = ""

        if not self.async_dump:
            cmd = [
                "adb", "-s", str(self.sn), "shell",
                "am start" + AppStartFactory.start_opt_dict[self.start_opt] + activity_name
            ]
            try:
                rlt_lines = execute_cmd_kill(cmd, self.sn, package_name)
                # 解析 WaitTime:xxx
                for raw in rlt_lines:
                    line = raw.strip()
                    if re.match(r"^WaitTime", line):
                        # 提取非字母和冒号的部分，即数值
                        m = re.search(r"[^a-zA-Z:]+", line)
                        if m:
                            wait_time_ms = int(m.group().strip())
                            break
                
                if wait_time_ms == 0 and rlt_lines:
                    error_msg = "未找到 WaitTime 或解析失败"
                    logger.debug(f"启动输出: {rlt_lines}")
                    
            except Exception as e:
                error_msg = str(e)
                logger.error(f"SerialNumber: {self.sn} --- 启动 {package_name} Exception!\n{e}")
                logger.error(traceback.format_exc())
                wait_time_ms = "Exception"

        end_time = time.time()
        elapsed_s = int(end_time - start_time + 1)
        return {
            "wait_time_ms": wait_time_ms,
            "elapsed_s": elapsed_s,
            "error": error_msg
        }


    def _respect_min_interval(self, logger, round_idx: int, package_name: str, elapsed_s: int):
        """确保任意两次启动之间至少间隔 app_start_remain 秒（与原实现一致）"""
        remain = max(0, int(getattr(self, "app_start_remain", 0)) - int(elapsed_s))
        if remain > 0:
            logger.info(
                f"SerialNumber: {self.sn} --- Round:{round_idx} --- 启动 {package_name} Done!, 等待{remain}s!\n"
            )
            time.sleep(remain)


    def _run_post_round_cmds(self, logger, round_idx: int):
        """每轮结束后执行命令并保存输出（与原路径/命名保持一致）"""
        for cmd_info in getattr(self, "cmd_list_after_round", []):
            if "###" not in cmd_info:
                continue
            shell_cmd, rlt_name = cmd_info.split("###", 1)
            cmdline = f"adb -s {self.sn} shell {shell_cmd}"
            rlt_file_path = os.path.join(self.result_dir, f"ROUND_{str(round_idx).zfill(4)}_{rlt_name}")
            logger.info(f"SerialNumber: {self.sn} --- Round:{round_idx}_ROUND_END --- 开始执行: {cmdline}")
            cmd_exec_and_save(cmdline, rlt_file_path, timeout=3600)


    def clear_flinger(self, max_retries=3):
        """安全清理 SurfaceFlinger，防止无限递归"""
        regex = re.compile(r"\\d+\\s+\\n(0\\s+0\\s+0\\s+\\n)+0\\s+0\\s+0$")
        return self._clear_flinger_recursive(max_retries, regex)
    
    def _clear_flinger_recursive(self, retries_left, regex):
        if retries_left <= 0:
            return False
            
        try:
            cmdline = f"adb -s {self.sn} shell dumpsys SurfaceFlinger --latency-clear"
            execute_cmd(cmdline)
            time.sleep(0.5)  # 等待清理完成
            
            cmdline = f"adb -s {self.sn} shell dumpsys SurfaceFlinger --latency"
            rlt = execute_cmd_rlt_read(cmdline)
            if isinstance(rlt, bytes):
                rlt = rlt.decode('utf-8', errors='ignore').strip()
            else:
                rlt = str(rlt).strip()
                
            if not re.match(regex, rlt):
                return self._clear_flinger_recursive(retries_left - 1, regex)
            return True
            
        except Exception as e:
            print(f"clear_flinger error: {e}")
            return False
    
    def _safe_remove_file(self, file_path):
        """安全删除文件"""
        try:
            if os.path.exists(file_path):
                os.remove(file_path)
                return True
        except Exception as e:
            print(f"删除文件失败 {file_path}: {e}")
            # 如果删除失败，尝试延迟后再次删除
            try:
                time.sleep(1)
                if os.path.exists(file_path):
                    os.remove(file_path)
            except:
                pass
        return False


if __name__ == "__main__":
    app = AppStartFactory("6465465465465465465", False, "xxx", "", "", "", False, "", "", 34, True)
    time.sleep(5)
    app.clear_flinger()
    time.sleep(10)